# -*- coding: utf-8 -*-

# system modules
import os
import json
import random
import time
import hashlib

# third party modules
import cv2
import numpy as np
from operator import itemgetter, attrgetter
from sklearn.cluster import KMeans
from sklearn import metrics

# current modules, it is necessary for this part
from . import tools


class ClassifySocketType(object):
    def __init__(self):
        """This class is used for classify and recognize socket
        Author
        ------
        caibingcheng

        Date
        ----
        2018.08.28

        Example
        -------
        Train:
            classify_socket_type = ClassifySocketType()
            try:
                classify_socket_type.train('./patches/', 'parameters', 'KMeans', 2)
            except Exception as e:
                print(e)
                exit(-1)

        Run:
            classify_socket_type = ClassifySocketType()
            try:
                classify_socket_type.setup('parameters')
            except Exception as e:
                print(e)
                exit(-1)
            ...#get frame
            classify_socket_type.run(frame)
            _, classify = classify_socket_type.getClassify()

        GetPatches:
            classify_socket_type = ClassifySocketType()
            # # optional
            # try:
            #     classify_socket_type.setup('parameters')
            # except Exception as e:
            #     print(e)
            #     exit(-1)
            try:
                classify_socket_type.getPatches(frame, './patches/')
            except Exception as e:  
                print(e)
                exit(-1)
        """

        self.m_contour_threshold = 1000
        self.m_area_threshold = 150000
        # for flexibleOTSU, default value
        self.m_classify_threshold = 100

        # container
        self.__m_center_select = []
        self.__m_contour_boxes_select = []
        self.__m_checking_boxes = []
        self.__m_classify_answer = []

        # file for save parameters, default value
        self.__m_save_patches_path = './patches/'

        # method for classifying/training/labeling, default value
        self.__m_classify_type = 'OTSU'
        self.__m_match_classify_type = {}

        # healthy flag, default value
        self.__m_run_success = False

        # patches size for sampling/training, default value
        self.m_roi_size = (8, 8)

        # method tuple, for checking correct inputing method.
        self.METHOD = ('OTSU', 'KMeans', )

        # wish you to add some Unsupervised Learning method here, for easily using.
        # method for training, add your training method here
        # you can reference the two methods listed here
        self.__m_trainMethod = {
            'OTSU': self.__trainOTSU,
            'KMeans': self.__trainKMeans
        }
        # classifier, add your classifier method here
        # you can reference the two methods listed here
        self.__m_classify = {
            'OTSU': self.__classifyOTSU,
            'KMeans': self.__classifyKMeans
        }
        # method for labeling, add your labeling method here
        # you can reference the two methods listed here
        # it is for Unsupervised Learning, which will make it more easier for using.
        self.__m_setLabel = {
            'OTSU': self.__setOUSTLabel,
            'KMeans': self.__setKMeansLabel
        }

    def __getContours(self, frame):
        """Get contours by binary image which is generated by OTSU

        Parameters
        ----------
        frame: the input image
            It will NOT be checked here.

        Return
        ------
        is_empty: If True, go on; Else, continue.
        contours: the contours
        """
        is_empty = True
        frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        ret, frame_bin = cv2.threshold(frame_gray, 0, 255, cv2.THRESH_OTSU)
        kernel = np.ones((15, 15), np.uint8)
        frame_bin = cv2.erode(frame_bin, kernel, iterations=2)
        _, contours_all, _ = cv2.findContours(
            frame_bin,
            cv2.RETR_EXTERNAL,
            cv2.CHAIN_APPROX_NONE)
        if contours_all.__len__() > 0:
            is_empty = False
        return is_empty, contours_all

    def __selectCounters(self, contours):
        """Select Contours by contours size (sum(pixels))

        Parameters
        ----------
        contours:
            The original contours returned by openCV

        Return
        ------
        is_empty: if all contours are sifted out, it returns False, else, True.
        boxes: the minmam bounding boxes of the remainder contours. 
               [(top-left, top-right, bottom-right, bottom-left), ]
        centers: the centers if the boxes. [((center point), (width, height)), ]
        """
        boxes, centers = [], []
        is_empty = True
        for contour in contours:
            if contour.size < self.m_contour_threshold:
                continue
            rect = cv2.minAreaRect(contour)
            rect = list(rect)
            rect[1] = list(rect[1])
            rect[1][0] *= 0.9
            rect[1][1] *= 0.9
            rect[1] = tuple(rect[1])
            rect = tuple(rect)
            box = cv2.boxPoints(rect)
            box = np.int0(box)
            boxes.append(box)
            centers.append((tuple(np.int0(rect[0])), rect[1]))
        if boxes.__len__() > 0:
            is_empty = False
        return is_empty, boxes, centers

    def __selectAreas(self, boxes, centers):
        """Select by area(u'面积')

        Parameters
        ----------
        boxes/centers: returned by __selectCounters()

        Return
        ------
        __m_center_select: the remainder centers(others are sifted out by its' small area)
        """
        is_empty = True
        for index, box in enumerate(boxes):
            area = cv2.contourArea(box)
            if area < self.m_area_threshold:
                continue
            self.__m_center_select.append(centers[index])
            self.__m_contour_boxes_select.append(box)
        if self.__m_center_select.__len__() > 0:
            is_empty = False
        return is_empty, self.__m_center_select

    def __generateCheckArea(self, centers):
        """Generated rect for checking

        Parameters
        ----------
        centers: returned by __selectAreas()

        Return
        ------
        checking_rect: The rects for classifying [((left-top point), (height, width)), ]
        """
        checking_rect = []
        is_empty = True
        for center in centers:
            line_len = min(center[1]) * 0.35  # len * 2^(1/2) / 2 / 2
            checking_box = [[(center[0][0] - line_len),
                             (center[0][1] - line_len)],
                            [(center[0][0] + line_len),
                             (center[0][1] - line_len)],
                            [(center[0][0] + line_len),
                             (center[0][1] + line_len)],
                            [(center[0][0] - line_len),
                             (center[0][1] + line_len)]]
            checking_box = np.int0(checking_box)
            checking_rect.append((tuple(checking_box[0]), tuple(
                np.int0([line_len * 2, line_len * 2]))))
            self.__m_checking_boxes.append(checking_box)
        if checking_rect.__len__() > 0:
            is_empty = False
        return is_empty, checking_rect

    def __getClassifyRect(self, frame):
        """Get rects for classifying

        Parameters
        ----------
        frame: image
            It will be checked here to ensure that it's not empty

        Return
        ------
        1: Correct flag, if this functions runs, it return True, else False
        2: The rects for classifying [((left-top point), (height, width)), ]
        3: The center and (height, width) [((center point), (height, width)), ]
        """
        empty, contours = self.__getContours(frame)
        if empty:
            return self.__m_run_success, None, None
        empty, boxes, centers = self.__selectCounters(contours)
        if empty:
            return self.__m_run_success, None, None
        empty, centers_select = self.__selectAreas(boxes, centers)
        if empty:
            return self.__m_run_success, None, None
        empty, classify_rect = self.__generateCheckArea(centers_select)
        if empty:
            return self.__m_run_success, None, None
        return True, classify_rect, centers_select

    def __loadParameter(self, train_file):
        """Load parameters from file in json format

        Parameters
        ----------
        train_file: string
            This file is usually generated by some ML method.
            The legality of this file path should be checked in the upper function

        Note
        ----
        The legality of json format and json attribute will be checked in
        this function.

        Exception Error
        ---------------
        Json Format Error: check the parameter file.
        Json Attribute Error: check the parameter file.
        """
        with open(train_file, 'r') as parameter_file:
            try:
                parameter = json.load(parameter_file)
            except:
                raise Exception('Json Format Error!')
            finally:
                parameter_file.close()
        try:
            self.__m_classify_type = parameter['classify_socket']['type']
            self.m_classify_threshold = parameter['classify_socket']['classify']
            self.m_roi_size = tuple(parameter['classify_socket']['roi'])
            label = tuple(parameter['classify_socket']['label'])
            self.m_contour_threshold = parameter['classify_socket']['contour_threshold']
            self.m_area_threshold = parameter['classify_socket']['area_threshold']
        except:
            raise Exception('Json Attribute Error!')

        if label.__len__() >= 2:
            self.setClassifyLabel(label)

    def __writeParameter(self, train_file, type, classify, label):
        """Write parameters to file

        Parameters
        ----------
        train_file: string
            The file which the parameters write to.
        classify: list[], single value
            The feature value of different classes.
            If it is generated by OUST, it's a single value.
            If it is generated by KMeans, it's a list value.
        label: list[]
            The label of each class, such as ['A', 'B',]

        Note
        ----
        If you change the attribute of this part, do not forget to update 
        those functions:
        __loadParameter()   setClassify()
        """
        with open(train_file, 'w') as parameter_file:
            data = {'classify_socket': {
                'type': type,
                'classify': classify,
                'roi': list(self.m_roi_size),
                'label': label,
                'contour_threshold': self.m_contour_threshold,
                'area_threshold': self.m_area_threshold}}
            json.dump(data, parameter_file, indent=4)
            parameter_file.close()

    def __trainOTSU(self, dir_path='./patches/', train_file='parameter', clusters=2):
        """Training classify threshold by OTSU

        Parameters
        ----------
        dir_path: string, optional
            The legality of this file path should be checked in the upper function.
            The training patches should be included in this directory.
        train_file: string, optional
            Containing the parameters of this algorithm.
        """
        self.__m_save_patches_path = dir_path
        image_dir = os.listdir(self.__m_save_patches_path)
        gray_histogram = []

        # get tainning patches
        for image_path in image_dir:
            if os.path.isdir(image_path):
                continue
            path = self.__m_save_patches_path + image_path
            image = cv2.imread(path)
            gray_value = np.mean(image)
            gray_histogram.append(gray_value)

        # if gray_histogram.__len__ <= 1, threshold = 100
        _, threshold = tools.flexibleOTSU(gray_histogram)
        self.__writeParameter(train_file, 'OTSU', threshold, ['A', 'B'])

    def __trainKMeans(self, dir_path='./patches/', train_file='parameter', clusters=2):
        """Training classify threshold by KMeans

        Parameters
        ----------
        dir_path: string, optional
            The legality of this file path should be checked in the upper function.
            The training patches should be included in this directory.
        train_file: string, optional
            Containing the parameters of this algorithm.

        Note
        ----
        sklearn required.

        Notice
        ------
        Notice the variable 'feature' and 'classify', which make it's possible to write
        different methods for training/classifying/labeling
        """
        self.__m_save_patches_path = dir_path
        image_dir = os.listdir(self.__m_save_patches_path)
        bgr_his = []

        # get training patches
        channels = (0, 1, 2)
        for image_path in image_dir:
            if os.path.isdir(image_path):
                continue
            path = self.__m_save_patches_path + image_path
            image = cv2.imread(path)
            sample = tools.getGivenChannelMean(image, channels)
            bgr_his.append(sample)

        # sklearn required
        kmeans = KMeans(clusters)
        kmeans.fit_predict(bgr_his)
        cluster_center = kmeans.cluster_centers_

        classify_mass = []
        for center in cluster_center:
            distancs_to_zero = center[0]**2 + center[1]**2 + center[2]**2
            classify_mass.append(
                # 'distance': the distance from the centre of gravity to original point
                # 'center': the centre of gravity of the class
                {'distance': distancs_to_zero,
                 'center': list(center)})
        feature = sorted(classify_mass, key=lambda x: x['distance'])
        # classify is a dict, so we can use a comman interface for different methods
        classify = {'feature': feature, 'channels': channels}
        # set default labels.
        label = [chr(0x41 + char) for char in range(0, clusters)]

        self.__writeParameter(train_file, 'KMeans', classify, label)

    def __classifyOTSU(self, patch, position):
        """Classify by OTSU

        Parameters
        ----------
        patch: image
            The roi of the checking area

        position: tuple() or list()
            The center point of the aim, which represent the aim.

        Note
        ----
        If you want to get the output of the classifier, call getClassify()
        """
        mean = np.mean(patch)
        classify = (mean <= self.m_classify_threshold)
        self.__m_classify_answer.append(
            (self.__m_match_classify_type[classify], position))

    def __classifyKMeans(self, patch, position):
        """Classify by KMeans

        Parameters
        ----------
        patch: image
            The roi of the checking area

        position: tuple() or list()
            The center point of the aim, which represent the aim.

        Note
        ----
        If you want to get the output of the classifier, call getClassify()
        """
        sample = tools.getGivenChannelMean(
            patch, self.m_classify_threshold['channels'])
        distance_mass = []
        for center in self.m_classify_threshold['feature']:
            distance = np.sqrt(
                np.sum(np.square(np.array(sample) -
                                 np.array((center['center'])))))
            distance_mass.append((distance, center['distance']))
        distance_sort = sorted(distance_mass, key=itemgetter(0, 1))
        self.__m_classify_answer.append(
            (self.__m_match_classify_type[distance_sort[0][1]], position))

    def __setKMeansLabel(self, classify_list):
        """Set the labels for KMeans

        Parameters
        ----------
        classify_list: list
            It should contains 2 labels at least, otherwise, return a Exception error.
            Labels are arranged by the given order, and the classes feature value are arranged
            from big to small by the distance to the original point, match them one by one.
            u'标签按照设定顺序，类别特征值按照离原点由大到小的顺序，依次匹配'

        Examples
        --------
        Given labels:
            ['X', 'Y', 'Z']
        The distance of the feature value to the original point:
            [5, 3, 4]
        Matched:
            'X' -- 5
            'Y' -- 4
            'Z' -- 3
        """
        distance_mass = []
        distance_sort = []
        for classify in self.m_classify_threshold['feature']:
            distance_mass.append(
                {'distance': classify['distance'], 'type': ''})
        distance_sort = sorted(distance_mass, key=lambda x: x['distance'])
        distance_sort.reverse()
        for index, classify_type in enumerate(classify_list):
            self.__m_match_classify_type[distance_sort[index]
                                         ['distance']] = classify_type

    def __setOUSTLabel(self, classify_list):
        """Set the labels for KMeans

        Note
        ----
        For OTSU method, if the feature value is greater than the threshold, it matches
        the first label. Otherwise, it matches the second label.
        """
        self.__m_match_classify_type[False] = classify_list[0]
        self.__m_match_classify_type[True] = classify_list[1]

    def selectROI(self, frame, rects):
        """Generate rois by rects

        Parameters
        ----------
        frame: image
        rects: list[] of rects
            The upper function has ensured the legality.

        Note
        ----
        Be Caution: No checking here.
        """
        rois = []
        for rect in rects:
            start_x = rect[0][1]
            end_x = start_x + rect[1][1]
            start_y = rect[0][0]
            end_y = start_y + rect[1][0]
            roi = frame[start_x:end_x, start_y:end_y]
            roi = cv2.resize(roi, self.m_roi_size)
            rois.append(roi)
        return rois

    def getPatches(self, frame, dir_path='./patches/'):
        """Get patches for trainning

        Parameters
        ----------
        frame: image
            If it is empty, return false
        dir_path: string, optional
            The directory of the trainning patches saved.

        Note
        ----
        As a public function, the legality of dir_path will be checked here.
        If the dir_path is not exist, generate one.
        """
        if not tools.checkEmptyData(frame):
            return False
        # generate the dir_path if it's not exist.
        tools.createEmptyDir(dir_path)
        # check again for safety
        tools.checkEmptyDir(dir_path)

        self.__m_save_patches_path = dir_path
        success, classify_rect, _ = self.__getClassifyRect(frame)
        if not success:
            return False
        rois = self.selectROI(frame, classify_rect)
        for roi in rois:
            # The patches file name is generated by the md5 code of the image.
            # The patch exactly the same with the existed one will not be saved.
            file_name = hashlib.md5(roi).hexdigest()
            cv2.imwrite(self.__m_save_patches_path + file_name + '.jpg', roi)
        self.clearBuffer()

        return True

    def train(self, dir_path='./patches/', train_file='parameter', method='OTSU', clusters=2):
        """Trainning function for public user.

        Parameters
        ----------
        dir_path: string, optional
            The directory contains the trainning patches.
            Will be checked here, if not exist, raise a Exception error.
        train_file: string, optional
            The file which the trained parameter write to.
            If exist, rewrite; if not exist, generate and write.
        method: string, optional
            Trainning method, it will be checked here.
            If it's not legal, raise a Exception error.

        Exception Error
        ---------------
        No Such Dir: check if is the dir_path correct
        No Such Method: ('OTSU', 'KMeans', ) only.
        """
        tools.checkEmptyDir(dir_path)
        tools.checkMethod(method, self.METHOD)
        # a dict contains the trainning methods
        self.__m_trainMethod[method](dir_path, train_file, clusters)

    def setup(self, train_file='parameter'):
        """Setup parameters by parameter_file

        Parameters
        ----------
        train_file: string, optional
            The parameter file for this class, it necessary.
            Will be checked, if not exist, raise a Exception error.

        Exception Error
        ---------------
        No Such File: check the parameter file, make sure it's exist.

        Note
        ----
        If this function is not be used, default parameters(OTSU) will be loaded.
        However, it is better if the parameters are under your control.
        """
        tools.checkEmptyFile(train_file)
        self.__loadParameter(train_file)

    def setClassifyLabel(self, classify_list=['A', 'B']):
        """Set the labels of differrent class

        Caution
        -------
        Not suggested. You'd better to load or modify labels by parameter file. See setup()

        Parameters
        ----------
        classify_list: list
            It should contains 2 labels at least, otherwise, return a Exception error.
            Labels are arranged by the given order, and the classes feature value are arranged
            from big to small by the distance to the original point, match them one by one.
            u'标签按照设定顺序，类别特征值按照离原点由大到小的顺序，依次匹配'

        Examples
        --------
        Given labels:
            ['X', 'Y', 'Z']
        The distance of the feature value to the original point:
            [5, 3, 4]
        Matched:
            'X' -- 5
            'Y' -- 4
            'Z' -- 3

        Note
        ----
        For OTSU method, if the feature value is greater than the threshold, it matches
        the first label. Otherwise, it matches the second label.

        Exception
        ---------
        Parameter Error: Check the sentence where calling this function
            or check the parameter file of its' label attribute.
        """
        tools.checkListLength(classify_list, 2)
        self.__m_setLabel[self.__m_classify_type](classify_list)

    def clearBuffer(self):
        """Clear the class buffer

        Example
        -------
        demo1:
            # Declare ClassifySocketType() outside loop
            classify_socket_type = ClassifySocketType()
            for frame in frame_set:
                classify_socket_type.run(frame)
                classify_socket_type.clearBuffer()

        demo2:
            # Declare ClassifySocketType() inside loop
            for frame in frame_set:
                classify_socket_type = ClassifySocketType()
                classify_socket_type.run(frame)

        For different situations, you are supposed to use differrent demo.

        Note
        ----
        Call this function in the end!!!
        In Example section, demo1 performs better.
        """
        self.__m_center_select = []
        self.__m_contour_boxes_select = []
        self.__m_checking_boxes = []
        self.__m_run_success = False
        self.__m_classify_answer = []

    def run(self, frame):
        """Run the algorithm

        Note
        ----
        The input frame will be check here to make sure it is not empty.
        If it runs, return true.
        If you get true, you are able to get the classify information by getClassify()
        """
        if not tools.checkEmptyData(frame):
            return False
        success, classify_rect, centers = self.__getClassifyRect(frame)
        if not success:
            return False
        # generate rois for classifying
        rois = self.selectROI(frame, classify_rect)
        for index, roi in enumerate(rois):
            self.__m_classify[self.__m_classify_type](roi, centers[index][0])
        self.__m_run_success = True
        return self.__m_run_success

    def getClassify(self):
        """Get the classify information

        Return
        ------
        False, None:
            If you get False, no aim detected.
        True, self.__m_classify_answer:
            If you get True, classify done.
            The second returned value [(Type, (Position, tuple)The Aim's Center Point), ].
            So, you can judge each aim by Type and it's Position.
            Otherwise, the Type is defined by setClassify() or modify the
            parameter file directly.

        Example
        -------
        If get True. You will get:
        [('A', (1938, 1347)), ('B', (2710, 1340))]
        [('A', (1921, 1348)), ('B', (2857, 1126))]
        [('A', (1454, 1481))]
        [('A', (1844, 1241))]

        If get False. You will get:
        None
        None
        None
        """
        if self.__m_run_success:
            return True, self.__m_classify_answer
        else:
            return False, None

    def draw(self, frame):
        """Draw the selecting rects & center points & type & classify roi

        Note
        ----
        Empty frame will be checked here.
        For debug.
        """
        if not tools.checkEmptyData(frame):
            return False
        draw_frame = frame.copy()

        if self.__m_run_success:
            for index, answer in enumerate(self.__m_classify_answer):
                # print(answer)
                cv2.putText(
                    draw_frame,
                    answer[0],
                    answer[1],
                    cv2.FONT_HERSHEY_SIMPLEX,
                    10,
                    (255, 0, 255),
                    5)
                cv2.circle(
                    draw_frame,
                    self.__m_center_select[index][0],
                    5,
                    (0, 255, 0),
                    10)
            cv2.drawContours(
                draw_frame, self.__m_contour_boxes_select, -1, (0, 0, 255), 3)
            cv2.drawContours(
                draw_frame, self.__m_checking_boxes, -1, (255, 0, 0), 3)

        show_frame = cv2.resize(draw_frame, (960, 640))
        cv2.imshow('classify', show_frame)
